import asyncio
import logging
import os

from llama_index.core import SimpleDirectoryReader
from llama_index.core import SummaryIndex
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.llms.dashscope import DashScope

logging.basicConfig(level=logging.DEBUG)

llm = DashScope(
    model="qwen-turbo",  # 使用通义千问Max模型 qwen-max
    api_key=os.getenv("DASHSCOPE_API_KEY"),  # 从阿里云控制台获取
    temperature=0.3
)

embed_model = HuggingFaceEmbedding(
    model_name="BAAI/bge-small-zh-v1.5",  # 中文优化模型
    device="cpu"
)


async def retrieve(retriever, query, label):
    response = await retriever.aretrieve(query)
    print(f"{label} retrieved {str(len(response))} nodes")


async def main():
    data_dir = 'D:/pythonRoom/llm_llama/data2'
    reader = SimpleDirectoryReader(input_dir=data_dir)
    documents = reader.load_data()
    index = SummaryIndex.from_documents(documents)
    retriever1 = index.as_retriever(
        retriever_mode='embedding',
        llm=llm,
        embed_model=embed_model
    )
    await asyncio.gather(
        retrieve(retriever1, '全职儿女', '<llm>'),
        retrieve(retriever1, '上班的好处', '<simple>')
    )


print('执行主函数')
asyncio.run(main())
